#include "storage/table/table_heap.h"

#include <cassert>

#include "common/logger.h"

TableHeap::TableHeap(BufferPoolManager *buffer_pool_manager, LockManager *lock_manager, LogManager *log_manager,
                     page_id_t first_page_id)
    : buffer_pool_manager_(buffer_pool_manager),
      lock_manager_(lock_manager),
      log_manager_(log_manager),
      first_page_id_(first_page_id) {}

TableHeap::TableHeap(BufferPoolManager *buffer_pool_manager, LockManager *lock_manager, LogManager *log_manager,
                     Transaction *txn)
    : buffer_pool_manager_(buffer_pool_manager), lock_manager_(lock_manager), log_manager_(log_manager) {
    // Initialize the first table page.
    auto first_page = reinterpret_cast<TablePage *>(buffer_pool_manager_->NewPage(&first_page_id_));
    DATABASE_ASSERT(first_page != nullptr, "Couldn't create a page for the table heap.");
    first_page->WLatch();
    first_page->Init(first_page_id_, PAGE_SIZE, INVALID_PAGE_ID, log_manager_, txn);
    first_page->WUnlatch();
    buffer_pool_manager_->UnpinPage(first_page_id_, true);
}

bool TableHeap::InsertTuple(const Tuple &tuple, RID *rid, Transaction *txn) {
    if (tuple.size_ + 32 > PAGE_SIZE) {  // larger than one page size
        txn->SetState(TransactionState::ABORTED);
        return false;
    }

    auto cur_page = static_cast<TablePage *>(buffer_pool_manager_->FetchPage(first_page_id_));
    if (cur_page == nullptr) {
        txn->SetState(TransactionState::ABORTED);
        return false;
    }

    cur_page->WLatch();
    // Insert into the first page with enough space. If no such page exists, create a new page and insert into that.
    // INVARIANT: cur_page is WLatched if you leave the loop normally.
    while (!cur_page->InsertTuple(tuple, rid, txn, lock_manager_, log_manager_)) {
        auto next_page_id = cur_page->GetNextPageId();
        // If the next page is a valid page,
        if (next_page_id != INVALID_PAGE_ID) {
            // Unlatch and unpin the current page.
            cur_page->WUnlatch();
            buffer_pool_manager_->UnpinPage(cur_page->GetTablePageId(), false);
            // And repeat the process with the next page.
            cur_page = static_cast<TablePage *>(buffer_pool_manager_->FetchPage(next_page_id));
            cur_page->WLatch();
        } else {
            // Otherwise we have run out of valid pages. We need to create a new page.
            auto new_page = static_cast<TablePage *>(buffer_pool_manager_->NewPage(&next_page_id));
            // If we could not create a new page,
            if (new_page == nullptr) {
                // Then life sucks and we abort the transaction.
                cur_page->WUnlatch();
                buffer_pool_manager_->UnpinPage(cur_page->GetTablePageId(), false);
                txn->SetState(TransactionState::ABORTED);
                return false;
            }
            // Otherwise we were able to create a new page. We initialize it now.
            new_page->WLatch();
            cur_page->SetNextPageId(next_page_id);
            new_page->Init(next_page_id, PAGE_SIZE, cur_page->GetTablePageId(), log_manager_, txn);
            cur_page->WUnlatch();
            buffer_pool_manager_->UnpinPage(cur_page->GetTablePageId(), true);
            cur_page = new_page;
        }
    }
    // This line has caused most of us to double-take and "whoa double unlatch".
    // We are not, in fact, double unlatching. See the invariant above.
    cur_page->WUnlatch();
    buffer_pool_manager_->UnpinPage(cur_page->GetTablePageId(), true);
    // Update the transaction's write set.
    txn->GetWriteSet()->emplace_back(*rid, WType::INSERT, Tuple{}, this);
    return true;
}

bool TableHeap::MarkDelete(const RID &rid, Transaction *txn) {
    // TODO(Amadou): remove empty page
    // Find the page which contains the tuple.
    auto page = reinterpret_cast<TablePage *>(buffer_pool_manager_->FetchPage(rid.GetPageId()));
    // If the page could not be found, then abort the transaction.
    if (page == nullptr) {
        txn->SetState(TransactionState::ABORTED);
        return false;
    }
    // Otherwise, mark the tuple as deleted.
    page->WLatch();
    page->MarkDelete(rid, txn, lock_manager_, log_manager_);
    page->WUnlatch();
    buffer_pool_manager_->UnpinPage(page->GetTablePageId(), true);
    // Update the transaction's write set.
    txn->GetWriteSet()->emplace_back(rid, WType::DELETE, Tuple{}, this);
    return true;
}

bool TableHeap::UpdateTuple(const Tuple &tuple, const RID &rid, Transaction *txn) {
    // Find the page which contains the tuple.
    auto page = reinterpret_cast<TablePage *>(buffer_pool_manager_->FetchPage(rid.GetPageId()));
    // If the page could not be found, then abort the transaction.
    if (page == nullptr) {
        txn->SetState(TransactionState::ABORTED);
        return false;
    }
    // Update the tuple; but first save the old value for rollbacks.
    Tuple old_tuple;
    page->WLatch();
    bool is_updated = page->UpdateTuple(tuple, &old_tuple, rid, txn, lock_manager_, log_manager_);
    page->WUnlatch();
    buffer_pool_manager_->UnpinPage(page->GetTablePageId(), is_updated);
    // Update the transaction's write set.
    if (is_updated && txn->GetState() != TransactionState::ABORTED) {
        txn->GetWriteSet()->emplace_back(rid, WType::UPDATE, old_tuple, this);
    }
    return is_updated;
}

void TableHeap::ApplyDelete(const RID &rid, Transaction *txn) {
    // Find the page which contains the tuple.
    auto page = reinterpret_cast<TablePage *>(buffer_pool_manager_->FetchPage(rid.GetPageId()));
    DATABASE_ASSERT(page != nullptr, "Couldn't find a page containing that RID.");
    // Delete the tuple from the page.
    page->WLatch();
    page->ApplyDelete(rid, txn, log_manager_);
    lock_manager_->Unlock(txn, rid);
    page->WUnlatch();
    buffer_pool_manager_->UnpinPage(page->GetTablePageId(), true);
}

void TableHeap::RollbackDelete(const RID &rid, Transaction *txn) {
    // Find the page which contains the tuple.
    auto page = reinterpret_cast<TablePage *>(buffer_pool_manager_->FetchPage(rid.GetPageId()));
    DATABASE_ASSERT(page != nullptr, "Couldn't find a page containing that RID.");
    // Rollback the delete.
    page->WLatch();
    page->RollbackDelete(rid, txn, log_manager_);
    page->WUnlatch();
    buffer_pool_manager_->UnpinPage(page->GetTablePageId(), true);
}

bool TableHeap::GetTuple(const RID &rid, Tuple *tuple, Transaction *txn) {
    // Find the page which contains the tuple.
    auto page = static_cast<TablePage *>(buffer_pool_manager_->FetchPage(rid.GetPageId()));
    // If the page could not be found, then abort the transaction.
    if (page == nullptr) {
        txn->SetState(TransactionState::ABORTED);
        return false;
    }
    // Read the tuple from the page.
    page->RLatch();
    bool res = page->GetTuple(rid, tuple, txn, lock_manager_);
    page->RUnlatch();
    buffer_pool_manager_->UnpinPage(rid.GetPageId(), false);
    return res;
}

void TableHeap::SetFirstPageId(page_id_t first_page_id) { first_page_id_ = first_page_id; }

void TableHeap::SetBufferPoolManager(BufferPoolManager *buffer_pool_manager) {
    buffer_pool_manager_ = buffer_pool_manager;
}

void TableHeap::SetLockManager(LockManager *lock_manager) { lock_manager_ = lock_manager; }

void TableHeap::SetLogManager(LogManager *log_manager) { log_manager_ = log_manager; }

TableIterator TableHeap::Begin(Transaction *txn) {
    // Start an iterator from the first page.
    // TODO(Wuwen): Hacky fix for now. Removing empty pages is a better way to handle this.
    RID rid;
    auto page_id = first_page_id_;
    while (page_id != INVALID_PAGE_ID) {
        auto page = static_cast<TablePage *>(buffer_pool_manager_->FetchPage(page_id));
        page->RLatch();
        // If this fails because there is no tuple, then RID will be the default-constructed value, which means EOF.
        auto found_tuple = page->GetFirstTupleRid(&rid);
        page->RUnlatch();
        buffer_pool_manager_->UnpinPage(page_id, false);
        if (found_tuple) {
            break;
        }
        page_id = page->GetNextPageId();
    }
    return TableIterator(this, rid, txn);
}

TableIterator TableHeap::End() { return TableIterator(this, RID(INVALID_PAGE_ID, 0), nullptr); }
